"""加密工具模块

提供区块链中间件系统中使用的各种加密功能，包括哈希、签名、加密解密等
"""

import hashlib
import hmac
import secrets
import base64
import json
from typing import Dict, Any, Optional, Tuple, Union
from datetime import datetime, timedelta
from cryptography.fernet import Fernet
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.backends import default_backend
import jwt
from .exceptions import CryptographyException, SignatureException, EncryptionException
from .logger import get_logger

logger = get_logger(__name__)

class HashUtils:
    """哈希工具类"""
    
    @staticmethod
    def sha256(data: Union[str, bytes]) -> str:
        """计算SHA256哈希"""
        if isinstance(data, str):
            data = data.encode('utf-8')
        return hashlib.sha256(data).hexdigest()
    
    @staticmethod
    def sha512(data: Union[str, bytes]) -> str:
        """计算SHA512哈希"""
        if isinstance(data, str):
            data = data.encode('utf-8')
        return hashlib.sha512(data).hexdigest()
    
    @staticmethod
    def md5(data: Union[str, bytes]) -> str:
        """计算MD5哈希（不推荐用于安全场景）"""
        if isinstance(data, str):
            data = data.encode('utf-8')
        return hashlib.md5(data).hexdigest()
    
    @staticmethod
    def hmac_sha256(data: Union[str, bytes], key: Union[str, bytes]) -> str:
        """计算HMAC-SHA256"""
        if isinstance(data, str):
            data = data.encode('utf-8')
        if isinstance(key, str):
            key = key.encode('utf-8')
        return hmac.new(key, data, hashlib.sha256).hexdigest()
    
    @staticmethod
    def keccak256(data: Union[str, bytes]) -> str:
        """计算Keccak256哈希（以太坊使用）"""
        try:
            from Crypto.Hash import keccak
            if isinstance(data, str):
                data = data.encode('utf-8')
            k = keccak.new(digest_bits=256)
            k.update(data)
            return k.hexdigest()
        except ImportError:
            logger.warning("pycryptodome not installed, falling back to SHA256")
            return HashUtils.sha256(data)
    
    @staticmethod
    def merkle_root(hashes: list) -> str:
        """计算Merkle树根哈希"""
        if not hashes:
            return HashUtils.sha256("")
        
        if len(hashes) == 1:
            return hashes[0]
        
        # 如果数量为奇数，复制最后一个哈希
        if len(hashes) % 2 == 1:
            hashes.append(hashes[-1])
        
        next_level = []
        for i in range(0, len(hashes), 2):
            combined = hashes[i] + hashes[i + 1]
            next_level.append(HashUtils.sha256(combined))
        
        return HashUtils.merkle_root(next_level)

class SymmetricEncryption:
    """对称加密工具类"""
    
    def __init__(self, key: Optional[bytes] = None):
        """初始化对称加密"""
        if key is None:
            key = Fernet.generate_key()
        elif isinstance(key, str):
            key = key.encode('utf-8')
        
        try:
            self.cipher = Fernet(key)
            self.key = key
        except Exception as e:
            raise EncryptionException(f"Failed to initialize symmetric encryption: {str(e)}")
    
    def encrypt(self, data: Union[str, bytes]) -> str:
        """加密数据"""
        try:
            if isinstance(data, str):
                data = data.encode('utf-8')
            encrypted = self.cipher.encrypt(data)
            return base64.b64encode(encrypted).decode('utf-8')
        except Exception as e:
            raise EncryptionException(f"Encryption failed: {str(e)}")
    
    def decrypt(self, encrypted_data: str) -> str:
        """解密数据"""
        try:
            encrypted_bytes = base64.b64decode(encrypted_data.encode('utf-8'))
            decrypted = self.cipher.decrypt(encrypted_bytes)
            return decrypted.decode('utf-8')
        except Exception as e:
            raise EncryptionException(f"Decryption failed: {str(e)}")
    
    def get_key(self) -> str:
        """获取密钥（Base64编码）"""
        return base64.b64encode(self.key).decode('utf-8')
    
    @staticmethod
    def generate_key() -> str:
        """生成新的密钥"""
        key = Fernet.generate_key()
        return base64.b64encode(key).decode('utf-8')

class AsymmetricEncryption:
    """非对称加密工具类"""
    
    def __init__(self, private_key: Optional[bytes] = None, public_key: Optional[bytes] = None):
        """初始化非对称加密"""
        try:
            if private_key is None and public_key is None:
                # 生成新的密钥对
                self.private_key = rsa.generate_private_key(
                    public_exponent=65537,
                    key_size=2048,
                    backend=default_backend()
                )
                self.public_key = self.private_key.public_key()
            else:
                if private_key:
                    self.private_key = serialization.load_pem_private_key(
                        private_key, password=None, backend=default_backend()
                    )
                    self.public_key = self.private_key.public_key()
                elif public_key:
                    self.public_key = serialization.load_pem_public_key(
                        public_key, backend=default_backend()
                    )
                    self.private_key = None
        except Exception as e:
            raise EncryptionException(f"Failed to initialize asymmetric encryption: {str(e)}")
    
    def encrypt(self, data: Union[str, bytes]) -> str:
        """使用公钥加密数据"""
        try:
            if isinstance(data, str):
                data = data.encode('utf-8')
            
            encrypted = self.public_key.encrypt(
                data,
                padding.OAEP(
                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
                    algorithm=hashes.SHA256(),
                    label=None
                )
            )
            return base64.b64encode(encrypted).decode('utf-8')
        except Exception as e:
            raise EncryptionException(f"Asymmetric encryption failed: {str(e)}")
    
    def decrypt(self, encrypted_data: str) -> str:
        """使用私钥解密数据"""
        if self.private_key is None:
            raise EncryptionException("Private key not available for decryption")
        
        try:
            encrypted_bytes = base64.b64decode(encrypted_data.encode('utf-8'))
            decrypted = self.private_key.decrypt(
                encrypted_bytes,
                padding.OAEP(
                    mgf=padding.MGF1(algorithm=hashes.SHA256()),
                    algorithm=hashes.SHA256(),
                    label=None
                )
            )
            return decrypted.decode('utf-8')
        except Exception as e:
            raise EncryptionException(f"Asymmetric decryption failed: {str(e)}")
    
    def sign(self, data: Union[str, bytes]) -> str:
        """使用私钥签名数据"""
        if self.private_key is None:
            raise SignatureException("Private key not available for signing")
        
        try:
            if isinstance(data, str):
                data = data.encode('utf-8')
            
            signature = self.private_key.sign(
                data,
                padding.PSS(
                    mgf=padding.MGF1(hashes.SHA256()),
                    salt_length=padding.PSS.MAX_LENGTH
                ),
                hashes.SHA256()
            )
            return base64.b64encode(signature).decode('utf-8')
        except Exception as e:
            raise SignatureException(f"Signing failed: {str(e)}")
    
    def verify(self, data: Union[str, bytes], signature: str) -> bool:
        """使用公钥验证签名"""
        try:
            if isinstance(data, str):
                data = data.encode('utf-8')
            
            signature_bytes = base64.b64decode(signature.encode('utf-8'))
            
            self.public_key.verify(
                signature_bytes,
                data,
                padding.PSS(
                    mgf=padding.MGF1(hashes.SHA256()),
                    salt_length=padding.PSS.MAX_LENGTH
                ),
                hashes.SHA256()
            )
            return True
        except Exception as e:
            logger.debug(f"Signature verification failed: {str(e)}")
            return False
    
    def get_private_key_pem(self) -> str:
        """获取私钥PEM格式"""
        if self.private_key is None:
            raise EncryptionException("Private key not available")
        
        pem = self.private_key.private_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PrivateFormat.PKCS8,
            encryption_algorithm=serialization.NoEncryption()
        )
        return pem.decode('utf-8')
    
    def get_public_key_pem(self) -> str:
        """获取公钥PEM格式"""
        pem = self.public_key.public_bytes(
            encoding=serialization.Encoding.PEM,
            format=serialization.PublicFormat.SubjectPublicKeyInfo
        )
        return pem.decode('utf-8')

class JWTManager:
    """JWT令牌管理器"""
    
    def __init__(self, secret_key: str, algorithm: str = 'HS256'):
        """初始化JWT管理器"""
        self.secret_key = secret_key
        self.algorithm = algorithm
    
    def generate_token(self, payload: Dict[str, Any], 
                      expires_in: int = 3600) -> str:
        """生成JWT令牌"""
        try:
            # 添加过期时间
            payload['exp'] = datetime.utcnow() + timedelta(seconds=expires_in)
            payload['iat'] = datetime.utcnow()
            
            token = jwt.encode(payload, self.secret_key, algorithm=self.algorithm)
            return token
        except Exception as e:
            raise CryptographyException(f"JWT token generation failed: {str(e)}")
    
    def verify_token(self, token: str) -> Dict[str, Any]:
        """验证JWT令牌"""
        try:
            payload = jwt.decode(token, self.secret_key, algorithms=[self.algorithm])
            return payload
        except jwt.ExpiredSignatureError:
            raise CryptographyException("JWT token has expired")
        except jwt.InvalidTokenError as e:
            raise CryptographyException(f"Invalid JWT token: {str(e)}")
    
    def refresh_token(self, token: str, expires_in: int = 3600) -> str:
        """刷新JWT令牌"""
        try:
            # 验证旧令牌（忽略过期）
            payload = jwt.decode(token, self.secret_key, 
                               algorithms=[self.algorithm], 
                               options={"verify_exp": False})
            
            # 移除时间相关字段
            payload.pop('exp', None)
            payload.pop('iat', None)
            
            # 生成新令牌
            return self.generate_token(payload, expires_in)
        except Exception as e:
            raise CryptographyException(f"JWT token refresh failed: {str(e)}")

class PasswordManager:
    """密码管理器"""
    
    @staticmethod
    def generate_salt() -> str:
        """生成盐值"""
        return secrets.token_hex(32)
    
    @staticmethod
    def hash_password(password: str, salt: str = None) -> Tuple[str, str]:
        """哈希密码"""
        if salt is None:
            salt = PasswordManager.generate_salt()
        
        # 使用PBKDF2进行密码哈希
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt.encode('utf-8'),
            iterations=100000,
            backend=default_backend()
        )
        
        password_hash = base64.b64encode(
            kdf.derive(password.encode('utf-8'))
        ).decode('utf-8')
        
        return password_hash, salt
    
    @staticmethod
    def verify_password(password: str, password_hash: str, salt: str) -> bool:
        """验证密码"""
        try:
            computed_hash, _ = PasswordManager.hash_password(password, salt)
            return hmac.compare_digest(password_hash, computed_hash)
        except Exception as e:
            logger.error(f"Password verification failed: {str(e)}")
            return False
    
    @staticmethod
    def generate_password(length: int = 16) -> str:
        """生成随机密码"""
        import string
        alphabet = string.ascii_letters + string.digits + "!@#$%^&*"
        return ''.join(secrets.choice(alphabet) for _ in range(length))

class APIKeyManager:
    """API密钥管理器"""
    
    @staticmethod
    def generate_api_key(length: int = 32) -> str:
        """生成API密钥"""
        return secrets.token_urlsafe(length)
    
    @staticmethod
    def generate_api_secret(length: int = 64) -> str:
        """生成API密钥"""
        return secrets.token_urlsafe(length)
    
    @staticmethod
    def create_api_signature(method: str, url: str, params: Dict[str, Any], 
                           secret: str, timestamp: str = None) -> str:
        """创建API签名"""
        if timestamp is None:
            timestamp = str(int(datetime.utcnow().timestamp()))
        
        # 构建签名字符串
        param_string = '&'.join([f"{k}={v}" for k, v in sorted(params.items())])
        sign_string = f"{method.upper()}&{url}&{param_string}&{timestamp}"
        
        # 计算HMAC签名
        signature = hmac.new(
            secret.encode('utf-8'),
            sign_string.encode('utf-8'),
            hashlib.sha256
        ).hexdigest()
        
        return signature
    
    @staticmethod
    def verify_api_signature(method: str, url: str, params: Dict[str, Any],
                           secret: str, signature: str, timestamp: str,
                           tolerance: int = 300) -> bool:
        """验证API签名"""
        try:
            # 检查时间戳
            current_time = int(datetime.utcnow().timestamp())
            request_time = int(timestamp)
            
            if abs(current_time - request_time) > tolerance:
                logger.warning(f"Request timestamp out of tolerance: {timestamp}")
                return False
            
            # 验证签名
            expected_signature = APIKeyManager.create_api_signature(
                method, url, params, secret, timestamp
            )
            
            return hmac.compare_digest(signature, expected_signature)
        except Exception as e:
            logger.error(f"API signature verification failed: {str(e)}")
            return False

class BlockchainCrypto:
    """区块链特定加密工具"""
    
    @staticmethod
    def generate_private_key() -> str:
        """生成私钥"""
        return secrets.token_hex(32)
    
    @staticmethod
    def private_key_to_public_key(private_key: str) -> str:
        """从私钥生成公钥（简化版）"""
        # 这里应该使用具体的椭圆曲线算法
        # 为了演示，使用哈希作为简化实现
        return HashUtils.sha256(private_key + "public")
    
    @staticmethod
    def public_key_to_address(public_key: str) -> str:
        """从公钥生成地址"""
        # 以太坊地址生成方式的简化版
        hash_result = HashUtils.keccak256(public_key)
        return "0x" + hash_result[-40:]
    
    @staticmethod
    def generate_wallet() -> Dict[str, str]:
        """生成钱包（私钥、公钥、地址）"""
        private_key = BlockchainCrypto.generate_private_key()
        public_key = BlockchainCrypto.private_key_to_public_key(private_key)
        address = BlockchainCrypto.public_key_to_address(public_key)
        
        return {
            'private_key': private_key,
            'public_key': public_key,
            'address': address
        }
    
    @staticmethod
    def sign_transaction(transaction_data: Dict[str, Any], private_key: str) -> str:
        """签名交易"""
        # 序列化交易数据
        tx_string = json.dumps(transaction_data, sort_keys=True)
        
        # 计算交易哈希
        tx_hash = HashUtils.sha256(tx_string)
        
        # 使用私钥签名（简化版）
        signature_data = tx_hash + private_key
        signature = HashUtils.sha256(signature_data)
        
        return signature
    
    @staticmethod
    def verify_transaction_signature(transaction_data: Dict[str, Any], 
                                   signature: str, public_key: str) -> bool:
        """验证交易签名"""
        try:
            # 重新计算交易哈希
            tx_string = json.dumps(transaction_data, sort_keys=True)
            tx_hash = HashUtils.sha256(tx_string)
            
            # 从公钥推导私钥哈希（这在实际中是不可能的，这里仅为演示）
            # 实际实现需要使用椭圆曲线数字签名算法
            expected_signature_data = tx_hash + HashUtils.sha256(public_key + "private")
            expected_signature = HashUtils.sha256(expected_signature_data)
            
            return hmac.compare_digest(signature, expected_signature)
        except Exception as e:
            logger.error(f"Transaction signature verification failed: {str(e)}")
            return False

# 工具函数
def generate_random_string(length: int = 32) -> str:
    """生成随机字符串"""
    return secrets.token_urlsafe(length)

def generate_uuid() -> str:
    """生成UUID"""
    import uuid
    return str(uuid.uuid4())

def constant_time_compare(a: str, b: str) -> bool:
    """常量时间字符串比较"""
    return hmac.compare_digest(a, b)

def secure_random_int(min_val: int = 0, max_val: int = 2**32) -> int:
    """生成安全随机整数"""
    return secrets.randbelow(max_val - min_val) + min_val

def encode_base64(data: Union[str, bytes]) -> str:
    """Base64编码"""
    if isinstance(data, str):
        data = data.encode('utf-8')
    return base64.b64encode(data).decode('utf-8')

def decode_base64(encoded_data: str) -> bytes:
    """Base64解码"""
    return base64.b64decode(encoded_data.encode('utf-8'))

def create_checksum(data: Union[str, bytes]) -> str:
    """创建校验和"""
    hash_value = HashUtils.sha256(data)
    return hash_value[:8]  # 取前8位作为校验和

def verify_checksum(data: Union[str, bytes], checksum: str) -> bool:
    """验证校验和"""
    expected_checksum = create_checksum(data)
    return constant_time_compare(checksum, expected_checksum)